package protocol

import (
	"encoding/binary"
	"io"
)

type CommandType uint32

const (
	READY CommandType = iota
	WAKE
	GIVE
	TAKE
	NO
	DONE
	STATUS // 新增
	CANCEL // 新增
)

type Message struct {
	Command CommandType
	Headers map[string]string
	Payload []byte
}

func (m *Message) Encode(w io.Writer) error {
	if err := binary.Write(w, binary.BigEndian, m.Command); err != nil {
		return err
	}

	headerCount := uint32(len(m.Headers))
	if err := binary.Write(w, binary.BigEndian, headerCount); err != nil {
		return err
	}

	for k, v := range m.Headers {
		if err := writeString(w, k); err != nil {
			return err
		}
		if err := writeString(w, v); err != nil {
			return err
		}
	}

	return writeBytes(w, m.Payload)
}

func (m *Message) Decode(r io.Reader) error {
	if err := binary.Read(r, binary.BigEndian, &m.Command); err != nil {
		return err
	}

	var headerCount uint32
	if err := binary.Read(r, binary.BigEndian, &headerCount); err != nil {
		return err
	}

	m.Headers = make(map[string]string)
	for i := uint32(0); i < headerCount; i++ {
		k, err := readString(r)
		if err != nil {
			return err
		}
		v, err := readString(r)
		if err != nil {
			return err
		}
		m.Headers[k] = v
	}

	payload, err := readBytes(r)
	if err != nil {
		return err
	}
	m.Payload = payload

	return nil
}

func writeString(w io.Writer, s string) error {
	if err := binary.Write(w, binary.BigEndian, uint32(len(s))); err != nil {
		return err
	}
	_, err := w.Write([]byte(s))
	return err
}

func readString(r io.Reader) (string, error) {
	var length uint32
	if err := binary.Read(r, binary.BigEndian, &length); err != nil {
		return "", err
	}
	buf := make([]byte, length)
	if _, err := io.ReadFull(r, buf); err != nil {
		return "", err
	}
	return string(buf), nil
}

func writeBytes(w io.Writer, b []byte) error {
	if err := binary.Write(w, binary.BigEndian, uint32(len(b))); err != nil {
		return err
	}
	_, err := w.Write(b)
	return err
}

func readBytes(r io.Reader) ([]byte, error) {
	var length uint32
	if err := binary.Read(r, binary.BigEndian, &length); err != nil {
		return nil, err
	}
	buf := make([]byte, length)
	if _, err := io.ReadFull(r, buf); err != nil {
		return nil, err
	}
	return buf, nil
}
